package com.example.test;

import org.openhitls.crypto.SM4;
import org.junit.Test;
import static org.junit.Assert.*;

import java.util.Arrays;

public class SM4Test {
    private static final byte[] TEST_KEY = new byte[] {
        (byte)0x01, (byte)0x02, (byte)0x03, (byte)0x04, (byte)0x05, (byte)0x06, (byte)0x07, (byte)0x08,
        (byte)0x09, (byte)0x0a, (byte)0x0b, (byte)0x0c, (byte)0x0d, (byte)0x0e, (byte)0x0f, (byte)0x10
    };

    private static final byte[] TEST_IV = new byte[] {
        (byte)0x11, (byte)0x22, (byte)0x33, (byte)0x44, (byte)0x55, (byte)0x66, (byte)0x77, (byte)0x88,
        (byte)0x99, (byte)0xaa, (byte)0xbb, (byte)0xcc, (byte)0xdd, (byte)0xee, (byte)0xff, (byte)0x00
    };

    private static final byte[] TEST_DATA = new byte[] {
        (byte)0x01, (byte)0x02, (byte)0x03, (byte)0x04, (byte)0x05, (byte)0x06, (byte)0x07, (byte)0x08,
        (byte)0x09, (byte)0x0a, (byte)0x0b, (byte)0x0c, (byte)0x0d, (byte)0x0e, (byte)0x0f, (byte)0x10,
        (byte)0x11, (byte)0x12, (byte)0x13, (byte)0x14, (byte)0x15, (byte)0x16, (byte)0x17, (byte)0x18,
        (byte)0x19, (byte)0x1a, (byte)0x1b, (byte)0x1c, (byte)0x1d, (byte)0x1e, (byte)0x1f, (byte)0x20
    };

    // Test data that's not block aligned (20 bytes)
    private static final byte[] UNALIGNED_DATA = new byte[] {
        (byte)0x01, (byte)0x02, (byte)0x03, (byte)0x04, (byte)0x05, (byte)0x06, (byte)0x07, (byte)0x08,
        (byte)0x09, (byte)0x0a, (byte)0x0b, (byte)0x0c, (byte)0x0d, (byte)0x0e, (byte)0x0f, (byte)0x10,
        (byte)0x11, (byte)0x12, (byte)0x13, (byte)0x14
    };

    private void printHex(String label, byte[] data) {
        System.out.print(label + ": ");
        for (byte b : data) {
            System.out.printf("%02x", b & 0xFF);
        }
        System.out.println();
    }

    @Test
    public void testSM4CBCPKCS7Padding() {
        System.out.println("\nRunning testSM4CBCPKCS7Padding");
        
        // Test with unaligned data to ensure padding works
        SM4 encryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_ENCRYPT, SM4.PADDING_PKCS7);
        SM4 decryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_DECRYPT, SM4.PADDING_PKCS7);

        printHex("Original data", UNALIGNED_DATA);

        // Encrypt
        byte[] encrypted = encryptor.encryptUpdate(UNALIGNED_DATA, 0, UNALIGNED_DATA.length);
        printHex("First update result", encrypted);
        byte[] finalEncBlock = encryptor.encryptFinal();
        printHex("Final encryption block", finalEncBlock);

        // Combine encrypted data
        byte[] fullEncrypted = combineArrays(encrypted, finalEncBlock);
        printHex("Full encrypted data", fullEncrypted);

        // Decrypt
        byte[] decrypted = decryptor.decryptUpdate(fullEncrypted, 0, fullEncrypted.length);
        printHex("First update result", decrypted);
        byte[] finalDecBlock = decryptor.decryptFinal(new int[1]);
        printHex("Final decryption block", finalDecBlock);

        // Combine decrypted data
        byte[] fullDecrypted = combineArrays(decrypted, finalDecBlock);
        printHex("Full decrypted data", fullDecrypted);

        assertArrayEquals("Decrypted data should match original data", UNALIGNED_DATA, fullDecrypted);
    }

    @Test
    public void testSM4CBCPKCS5Padding() {
        System.out.println("\nRunning testSM4CBCPKCS5Padding");
        
        // Test with unaligned data to ensure padding works
        SM4 encryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_ENCRYPT, SM4.PADDING_PKCS5);
        SM4 decryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_DECRYPT, SM4.PADDING_PKCS5);

        printHex("Original data", UNALIGNED_DATA);

        // Encrypt
        byte[] encrypted = encryptor.encryptUpdate(UNALIGNED_DATA, 0, UNALIGNED_DATA.length);
        printHex("First update result", encrypted);
        byte[] finalEncBlock = encryptor.encryptFinal();
        printHex("Final encryption block", finalEncBlock);

        // Combine encrypted data
        byte[] fullEncrypted = combineArrays(encrypted, finalEncBlock);
        printHex("Full encrypted data", fullEncrypted);

        // Decrypt
        byte[] decrypted = decryptor.decryptUpdate(fullEncrypted, 0, fullEncrypted.length);
        printHex("First update result", decrypted);
        byte[] finalDecBlock = decryptor.decryptFinal(new int[1]);
        printHex("Final decryption block", finalDecBlock);

        // Combine decrypted data
        byte[] fullDecrypted = combineArrays(decrypted, finalDecBlock);
        printHex("Full decrypted data", fullDecrypted);

        assertArrayEquals("Decrypted data should match original data", UNALIGNED_DATA, fullDecrypted);
    }

    @Test
    public void testSM4CBCZerosPadding() {
        System.out.println("\nRunning testSM4CBCZerosPadding");
        
        // Test with unaligned data to ensure padding works
        SM4 encryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_ENCRYPT, SM4.PADDING_ZEROS);
        SM4 decryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_DECRYPT, SM4.PADDING_ZEROS);

        printHex("Original data", UNALIGNED_DATA);

        // Encrypt
        byte[] encrypted = encryptor.encryptUpdate(UNALIGNED_DATA, 0, UNALIGNED_DATA.length);
        printHex("First update result", encrypted);
        byte[] finalEncBlock = encryptor.encryptFinal();
        printHex("Final encryption block", finalEncBlock);

        // Combine encrypted data
        byte[] fullEncrypted = combineArrays(encrypted, finalEncBlock);
        printHex("Full encrypted data", fullEncrypted);

        // Decrypt
        byte[] decrypted = decryptor.decryptUpdate(fullEncrypted, 0, fullEncrypted.length);
        printHex("First update result", decrypted);
        byte[] finalDecBlock = decryptor.decryptFinal(new int[1]);
        printHex("Final decryption block", finalDecBlock);

        // Combine decrypted data
        byte[] fullDecrypted = combineArrays(decrypted, finalDecBlock);
        printHex("Full decrypted data", fullDecrypted);

        // For zeros padding, we need to trim trailing zeros
        int actualLength = UNALIGNED_DATA.length;
        byte[] trimmedDecrypted = Arrays.copyOf(fullDecrypted, actualLength);
        assertArrayEquals("Decrypted data should match original data", UNALIGNED_DATA, trimmedDecrypted);
    }

    @Test
    public void testSM4CBCISO7816Padding() {
        System.out.println("\nRunning testSM4CBCISO7816Padding");
        
        // Test with unaligned data to ensure padding works
        SM4 encryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_ENCRYPT, SM4.PADDING_ISO7816);
        SM4 decryptor = new SM4(SM4.SM4_CBC, TEST_KEY, TEST_IV, SM4.MODE_DECRYPT, SM4.PADDING_ISO7816);

        printHex("Original data", UNALIGNED_DATA);

        // Encrypt
        byte[] encrypted = encryptor.encryptUpdate(UNALIGNED_DATA, 0, UNALIGNED_DATA.length);
        printHex("First update result", encrypted);
        byte[] finalEncBlock = encryptor.encryptFinal();
        printHex("Final encryption block", finalEncBlock);

        // Combine encrypted data
        byte[] fullEncrypted = combineArrays(encrypted, finalEncBlock);
        printHex("Full encrypted data", fullEncrypted);

        // Decrypt
        byte[] decrypted = decryptor.decryptUpdate(fullEncrypted, 0, fullEncrypted.length);
        printHex("First update result", decrypted);
        byte[] finalDecBlock = decryptor.decryptFinal(new int[1]);
        printHex("Final decryption block", finalDecBlock);

        // Combine decrypted data
        byte[] fullDecrypted = combineArrays(decrypted, finalDecBlock);
        printHex("Full decrypted data", fullDecrypted);

        assertArrayEquals("Decrypted data should match original data", UNALIGNED_DATA, fullDecrypted);
    }

    private byte[] combineArrays(byte[] first, byte[] second) {
        if (first == null) return second;
        if (second == null) return first;
        byte[] result = new byte[first.length + second.length];
        System.arraycopy(first, 0, result, 0, first.length);
        System.arraycopy(second, 0, result, first.length, second.length);
        return result;
    }
} 